import numpy as np
import warnings
import os

def gen_golden_data_simple():
    input_shape = [7, 117]
    output_shape = [7, 117]

    warnings.filterwarnings('ignore', category=RuntimeWarning)
    def float32_to_bf16(x):
        x32 = x.astype(np.float32)
        u32 = x32.view(np.uint32)
        bf16 = (u32 >> 16).astype(np.uint16)
        return bf16

    def bf16_to_float32(bf16):
        u32 = (bf16.astype(np.uint32)) << 16
        return u32.view(np.float32)
    

    testid = 8  # 你可以在这里改测试编号

    print(f"-------------TEST {testid}--------------")

    def handle_negative_exp(x, y):
        """负指数时取底数倒数，防止溢出/NaN"""
        base = np.where(y < 0, 1.0 / (np.abs(x) + 1e-6), x)
        return base, np.abs(y)

    match testid:
        # ==============================
        # case 0: float32 , float32 -> float32
        # ==============================
        case 0:
            x = np.random.uniform(0, 10, input_shape).astype(np.float32)
            y = np.random.uniform(0, 3, input_shape).astype(np.float32)
            base, y_abs = handle_negative_exp(x, y)
            golden = (x**y).astype(np.float32)
            print("[INFO] 已生成: float32 输入, float32 输入, float32 输出")
            
        # ==============================
        # case 1: uint8 , uint8 -> uint8
        # ==============================
        case 1:
            x = np.random.randint(0, 6, size=input_shape, dtype=np.uint8)
            y = np.random.randint(0, 3, size=input_shape, dtype=np.uint8)
            golden = (x ** y).astype(np.uint8)
            print("[INFO] 已生成: uint8 输入, uint8 输入, uint8 输出")

         # ==============================
        # case 2: int32 , int32 -> int32
        # ==============================
        case 2:
            x = np.random.randint(0, 6, size=input_shape, dtype=np.int32)
            y = np.random.randint(0, 5, size=input_shape, dtype=np.int32)
            base, y_abs = handle_negative_exp(x, y)
            golden = (base ** y_abs).astype(np.int32)
            print("[INFO] 已生成: int32 输入, int32 输入, int32 输出")

         # ==============================
        # case 3: 广播
        # ==============================
        case 3:
            x_shape = [4,1,5]
            y_shape = [1,3,1]
            x = np.random.randint(4, 6, size = x_shape, dtype=np.uint8)
            y = np.random.randint(0, 3, size = y_shape, dtype=np.uint8)
            x_brc, y_brc = np.broadcast_arrays(x, y)
            x_brc.tofile("./input/input_x_brc.bin")
            y_brc.tofile("./input/input_y_brc.bin")
            golden = (x_brc.astype(np.uint8) ** y_brc.astype(np.uint8)).astype(np.uint8)

            print("[INFO] 已生成: uint8 输入, uint8 输入, uint8 输出")

        
         # ==============================
        # case 4: int16, int16 ->int 16
        # ==============================
        case 4:
            x = np.random.randint(0, 10, size=input_shape, dtype=np.int16)
            y = np.random.randint(0, 5, size=input_shape, dtype=np.int16)
            base, y_abs = handle_negative_exp(x, y)
            golden = (base ** y_abs).astype(np.int16)
            print("[INFO] 已生成: int16 输入, int16 输入, int16 输出")
        
        # ==============================
        # case 5: int8, int8 ->int8
        # ==============================
        case 5:
            x = np.random.randint(0, 5, size=input_shape, dtype=np.int8)
            y = np.random.randint(0, 3, size=input_shape, dtype=np.int8)
            base, y_abs = handle_negative_exp(x, y)
            golden = (base ** y_abs).astype(np.int8)
            print("[INFO] 已生成: int8 输入, int8 输入, int8 输出")

        # ==============================
        # case 6: float16, float16 ->float16
        # ==============================
        case 6:
            x = np.random.uniform(0, 10, input_shape).astype(np.float16)
            y = np.random.uniform(0, 3, input_shape).astype(np.float16)

            base, y_abs = handle_negative_exp(x, y)
            golden = (base ** y_abs).astype(np.float16)
            print("[INFO] 已生成: float16 输入, float16 输入, float16 输出")
        
        # ==============================
        # case 7: bfloat16, bfloat16 ->bfloat16
        # ==============================
        case 7:
            x_f32 = np.random.uniform(0, 10, input_shape).astype(np.float32)
            y_f32 = np.random.uniform(0, 3, input_shape).astype(np.float32)
            x = float32_to_bf16(x_f32)
            y = float32_to_bf16(y_f32)
            x32 = bf16_to_float32(x)
            y32 = bf16_to_float32(y)
            golden32 = x32 ** y32
            golden = float32_to_bf16(golden32)

            print("[INFO] 已生成: bfloat16 输入, bfloat16 输入, bfloat16 输出")
            
        # ==============================
        # case 8: int16, float32 ->float32
        # ==============================
        case 8:
            x = np.random.randint(0, 5, size=input_shape, dtype=np.int16)
            y = np.random.uniform(0, 4, input_shape).astype(np.float32)
            x_32 = x.astype(np.float32)
            base, y_abs = handle_negative_exp(x_32, y)
            golden = (base ** y_abs).astype(np.float32)

            print("[INFO] 已生成: int16 输入, float32 输入, float32 输出")

        # ==============================
        # case 9: int8, float32 ->float32
        # ==============================
        case 9:
            x = np.random.randint(0, 5, size=input_shape, dtype=np.int8)
            y = np.random.uniform(0, 4, input_shape).astype(np.float32)
            x_32 = x.astype(np.float32)
            base, y_abs = handle_negative_exp(x_32, y)
            golden = (base ** y_abs).astype(np.float32)

            print("[INFO] 已生成: int8 输入, float32 输入, float32 输出")



        # ==============================
        # case 20: 充分广播（fully broadcasted tensor）
        # ==============================
        case 20:
            # x1 形状 [2,1,9]
            x = np.random.uniform(0, 4, (2,1,9)).astype(np.float32)  
            # x2 形状 [1,4,1]
            y = np.random.uniform(1, 5, (1,4,1)).astype(np.float32)  
            # 输出 golden 使用 numpy 广播规则
            golden = (x ** y).astype(np.float32)

            print("[INFO] 已生成: float32 张量输入, float32 张量输入, float32 输出")
            print("x.shape =", x.shape, ", y.shape =", y.shape, ", golden.shape =", golden.shape)
        # ==============================
        # 默认
        # ==============================
        case _:
            print("未匹配到")
            return

    os.system("mkdir -p input")
    os.system("mkdir -p output")

    x.tofile("./input/input_x.bin")
    y.tofile("./input/input_y.bin")
    golden.tofile("./output/golden.bin")
    print("[SAVED] 已保存输入与golden输出。")

if __name__ == "__main__":
    gen_golden_data_simple()
